# -*- coding: utf-8 -*-
# @Time    : 2025/1/9 下午1:54
# @Author  : ysj
# @FileName: idl2txt.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/ydscc?type=blog
import os
import cv2
import shutil
"""
将brainwash_train.idl文件中的标注数据转换为YOLO格式。

Args:
    idl_file_path (str): brainwash_train.idl 文件的路径。
    images_folder (str): 图片文件夹路径。
    output_folder (str): YOLO 格式标注保存文件夹路径。
"""
def convert_idl_to_yolo(idl_file_path, images_folder, output_folder):
    """
    将brainwash_train.idl文件中的标注数据转换为YOLO格式。

    Args:
        idl_file_path (str): brainwash_train.idl 文件的路径。
        images_folder (str): 图片文件夹路径。
        output_folder (str): YOLO 格式标注保存文件夹路径。
    """
    os.makedirs(output_folder, exist_ok=True)

    # 打开并读取IDL文件
    with open(idl_file_path, 'r') as file:
        lines = file.readlines()

    for line in lines:
        # 跳过空行或无效行
        if not line.strip() or ':' not in line:
            continue

        # 提取图像路径和标注信息
        if line.strip().split(":"):
            image_path, bboxes_str = line.strip().split(":")
            image_path = image_path.strip().strip('"')
            bboxes_str = bboxes_str.strip().rstrip(';')
        else:
            continue

        # 解析每个边界框
        bboxes = []
        # 去除多余空格并正确拆分
        for bbox in bboxes_str.split("),"):
            bbox = bbox.strip('()').strip()
            if bbox:  # 忽略空字符串
                # 修正并拆分字符串
                bbox_parts = bbox.replace('(', '').replace(')', '').split(',')
                # 确保转换为浮点数
                try:
                    bboxes.append(tuple(map(float, bbox_parts)))
                except ValueError:
                    print(f"Skipping invalid bbox: {bbox}")
                    continue

        # 获取图像路径和完整路径
        full_image_path = os.path.join(images_folder, image_path)

        if not os.path.exists(full_image_path):
            print(f"Warning: Image {full_image_path} does not exist. Skipping...")
            continue

        # 读取图像并获取尺寸
        img = cv2.imread(full_image_path)
        if img is None:
            print(f"Warning: Could not read image {full_image_path}. Skipping...")
            continue
        img_h, img_w = img.shape[:2]
        res_img = os.path.join(output_folder, os.path.basename(image_path))
        # 创建对应的YOLO标注文件
        txt_output_path = os.path.join(output_folder, os.path.splitext(os.path.basename(image_path))[0] + '.txt')
        if os.path.exists(txt_output_path):
            txt_output_path = os.path.join(output_folder, os.path.splitext(os.path.basename(image_path))[0]+f'_{1}.txt')
            res_img = os.path.join(output_folder, os.path.basename(image_path).replace('.', '_1.'))
        shutil.move(full_image_path, res_img)
        with open(txt_output_path, 'w') as txt_file:
            # 解析每个边界框并转换为YOLO格式
            for bbox in bboxes:
                xmin, ymin, xmax, ymax = bbox
                x_center = ((xmin + xmax) / 2) / img_w
                y_center = ((ymin + ymax) / 2) / img_h
                width = (xmax - xmin) / img_w
                height = (ymax - ymin) / img_h

                # 假设类别ID为0（即人头检测）
                txt_file.write(f"0 {x_center} {y_center} {width} {height}\n")

        print(f"Converted: {image_path} -> {txt_output_path}")


def find_path():
    from collections import Counter
    folder_path1 = r'C:\Al\Software\AI_Model\Project\yolov5\yolov5\train_data\train\images'  # 替换为实际的文件夹路径
    folder_path2 = r'C:\Al\Software\AI_Model\Project\yolov5\yolov5\train_data\train\labels'
    images_folder = []
    labels_folder = []
    # 获取文件夹中的所有文件
    files = os.listdir(folder_path1)
    label = os.listdir(folder_path2)

    # 筛选包含 "_1" 的文件
    files_with_1 = [file for file in files if '_1' in file]
    files_with_2 = [file for file in label if '_1' in file]
    # 输出结果
    for file in files_with_1:
        images_folder.append(os.path.splitext(file)[0])
    for file in files_with_2:
        labels_folder.append(os.path.splitext(file)[0])


    diff1 = list((Counter(images_folder) - Counter(labels_folder)).elements())
    print(len(diff1))
    print(f"List1 unique elements (with frequency): {diff1}")


def del_dif():
    import os

    # 给定的文件名列表
    file_names_to_delete = [
        '00454000_640x480_1', '00455000_640x480_1', '00456000_640x480_1', '00457000_640x480_1',
        '00458000_640x480_1', '00459000_640x480_1', '00460000_640x480_1', '00461000_640x480_1',
        '00462000_640x480_1', '00463000_640x480_1', '00464000_640x480_1', '00465000_640x480_1',
        '00466000_640x480_1', '00468000_640x480_1', '00469000_640x480_1', '00471000_640x480_1',
        '00473000_640x480_1', '00474000_640x480_1', '00475000_640x480_1', '00477000_640x480_1',
        '00478000_640x480_1', '00479000_640x480_1', '00480000_640x480_1', '00481000_640x480_1'
    ]

    # 指定目标文件夹路径
    folder_path = r"C:\Al\Software\AI_Model\Project\yolov5\yolov5\train_data\train"

    # 遍历文件夹中的所有文件
    for filename in os.listdir(folder_path):
        # 检查文件名是否包含在给定的列表中
        if any(file_name in filename for file_name in file_names_to_delete):
            file_path = os.path.join(folder_path, filename)
            try:
                # 删除文件
                os.remove(file_path)
                print(f"Deleted: {file_path}")
            except Exception as e:
                print(f"Error deleting {file_path}: {e}")


# # 使用示例
# idl_path = r"C:\Al\Software\AI_Model\Project\yolov5\yolov5\my_data\person_face\brainwash\brainwash_val.idl"  # Brainwash 数据集的 .idl 文件路径
# images_folder = r"C:\Al\Software\AI_Model\Project\yolov5\yolov5\my_data\person_face\brainwash"  # 图片文件夹路径
# yolo_output_folder = r"C:\Al\Software\AI_Model\Project\yolov5\yolov5\my_data\person_face\brainwash\val_label"  # YOLO 标注保存文件夹路径
#
# convert_idl_to_yolo(idl_path, images_folder, yolo_output_folder)
find_path()

